/**
 * Loader for KTX 2.0 GPU Texture containers.
 *
 * KTX 2.0 is a container format for various GPU texture formats. The loader
 * supports Basis Universal GPU textures, which can be quickly transcoded to
 * a wide variety of GPU texture compression formats, as well as some
 * uncompressed DataTexture and Data3DTexture formats.
 *
 * References:
 * - KTX: http://github.khronos.org/KTX-Specification/
 * - DFD: https://www.khronos.org/registry/DataFormat/specs/1.3/dataformat.1.3.html#basicdescriptor
 */

import {
  CompressedTexture,
  CompressedArrayTexture,
  Data3DTexture,
  DataTexture,
  FileLoader,
  FloatType,
  HalfFloatType,
  NoColorSpace,
  LinearFilter,
  LinearMipmapLinearFilter,
  Loader,
  RedFormat,
  RGB_ETC1_Format,
  RGB_ETC2_Format,
  RGB_PVRTC_4BPPV1_Format,
  RGB_S3TC_DXT1_Format,
  RGBA_ASTC_4x4_Format,
  RGBA_BPTC_Format,
  RGBA_ETC2_EAC_Format,
  RGBA_PVRTC_4BPPV1_Format,
  RGBA_S3TC_DXT5_Format,
  RGBAFormat,
  RGFormat,
  SRGBColorSpace,
  UnsignedByteType,
} from "three";
import { WorkerPool } from "three/examples/jsm/utils/WorkerPool";
import {
  read,
  KHR_DF_FLAG_ALPHA_PREMULTIPLIED,
  KHR_DF_TRANSFER_SRGB,
  KHR_SUPERCOMPRESSION_NONE,
  KHR_SUPERCOMPRESSION_ZSTD,
  VK_FORMAT_UNDEFINED,
  VK_FORMAT_R16_SFLOAT,
  VK_FORMAT_R16G16_SFLOAT,
  VK_FORMAT_R16G16B16A16_SFLOAT,
  VK_FORMAT_R32_SFLOAT,
  VK_FORMAT_R32G32_SFLOAT,
  VK_FORMAT_R32G32B32A32_SFLOAT,
  VK_FORMAT_R8_SRGB,
  VK_FORMAT_R8_UNORM,
  VK_FORMAT_R8G8_SRGB,
  VK_FORMAT_R8G8_UNORM,
  VK_FORMAT_R8G8B8A8_SRGB,
  VK_FORMAT_R8G8B8A8_UNORM,
} from "three/examples/jsm/libs/ktx-parse.module";
import { ZSTDDecoder } from "three/examples/jsm/libs/zstddec.module";

const _taskCache = new WeakMap();

let _activeLoaders = 0;

let _zstd;

class KTX2Loader extends Loader {
  constructor(manager) {
    super(manager);

    this.transcoderPath = "";
    this.transcoderBinary = null;
    this.transcoderPending = null;

    this.workerPool = new WorkerPool();
    this.workerSourceURL = "";
    this.workerConfig = null;

    if (typeof MSC_TRANSCODER !== "undefined") {
      console.warn(
        'THREE.KTX2Loader: Please update to latest "basis_transcoder".' +
          ' "msc_basis_transcoder" is no longer supported in three.js r125+.'
      );
    }
  }

  setTranscoderPath(path) {
    this.transcoderPath = path;

    return this;
  }

  setWorkerLimit(num) {
    this.workerPool.setWorkerLimit(num);

    return this;
  }

  detectSupport(renderer) {
    if (renderer.isWebGPURenderer === true) {
      this.workerConfig = {
        astcSupported: renderer.hasFeature("texture-compression-astc"),
        etc1Supported: false,
        etc2Supported: renderer.hasFeature("texture-compression-etc2"),
        dxtSupported: renderer.hasFeature("texture-compression-bc"),
        bptcSupported: false,
        pvrtcSupported: false,
      };
    } else {
      this.workerConfig = {
        astcSupported: renderer.extensions.has("WEBGL_compressed_texture_astc"),
        etc1Supported: renderer.extensions.has("WEBGL_compressed_texture_etc1"),
        etc2Supported: renderer.extensions.has("WEBGL_compressed_texture_etc"),
        dxtSupported: renderer.extensions.has("WEBGL_compressed_texture_s3tc"),
        bptcSupported: renderer.extensions.has("EXT_texture_compression_bptc"),
        pvrtcSupported:
          renderer.extensions.has("WEBGL_compressed_texture_pvrtc") ||
          renderer.extensions.has("WEBKIT_WEBGL_compressed_texture_pvrtc"),
      };

      if (renderer.capabilities.isWebGL2) {
        // https://github.com/mrdoob/three.js/pull/22928
        this.workerConfig.etc1Supported = false;
      }
    }

    return this;
  }

  init() {
    if (!this.transcoderPending) {
      this.transcoderPending = Promise.all([]).then(([]) => {
        this.workerPool.setWorkerCreator(() => {
          const worker = new Worker("workers/Ktx2.js");
          const transcoderBinary = new Uint8Array([]);

          worker.postMessage({ type: "init", config: this.workerConfig, transcoderBinary }, [transcoderBinary]);

          return worker;
        });
      });

      if (_activeLoaders > 0) {
        // Each instance loads a transcoder and allocates workers, increasing network and memory cost.

        console.warn(
          "THREE.KTX2Loader: Multiple active KTX2 loaders may cause performance issues." +
            " Use a single KTX2Loader instance, or call .dispose() on old instances."
        );
      }

      _activeLoaders++;
    }

    return this.transcoderPending;
  }

  loadBuffer(buffer, onLoad, onProgress, onError) {
    if (_taskCache.has(buffer)) {
      const cachedTask = _taskCache.get(buffer);
      return cachedTask.promise.then(onLoad).catch(onError);
    }

    this._createTexture(buffer)
      .then((texture) => (onLoad ? onLoad(texture) : null))
      .catch(onError);
  }

  load(url, onLoad, onProgress, onError) {
    if (this.workerConfig === null) {
      throw new Error("THREE.KTX2Loader: Missing initialization with `.detectSupport( renderer )`.");
    }

    const loader = new FileLoader(this.manager);

    loader.setResponseType("arraybuffer");
    loader.setWithCredentials(this.withCredentials);

    loader.load(
      url,
      (buffer) => {
        // Check for an existing task using this buffer. A transferred buffer cannot be transferred
        // again from this thread.
        if (_taskCache.has(buffer)) {
          const cachedTask = _taskCache.get(buffer);

          return cachedTask.promise.then(onLoad).catch(onError);
        }

        this._createTexture(buffer)
          .then((texture) => (onLoad ? onLoad(texture) : null))
          .catch(onError);
      },
      onProgress,
      onError
    );
  }

  _createTextureFrom(transcodeResult, container) {
    const { mipmaps, width, height, format, type, error, dfdTransferFn, dfdFlags } = transcodeResult;
    for (let i = 0; i < transcodeResult.mipmaps.length; i++) {
      transcodeResult.mipmaps[i].data = new Uint8Array(transcodeResult.buffers[i]);
      //   debugger
      //   let oldData = m.data;
      //   m.data = [];
      //   for (const k in oldData) {
      //     m.data[k] = oldData[k];
      //   }
      //   m.data = new Uint8Array(m.data);
      //   oldData = null;
    }

    if (type === "error") return Promise.reject(error);

    const texture =
      container.layerCount > 1
        ? new CompressedArrayTexture(mipmaps, width, height, container.layerCount, format, UnsignedByteType)
        : new CompressedTexture(mipmaps, width, height, format, UnsignedByteType);

    texture.minFilter = mipmaps.length === 1 ? LinearFilter : LinearMipmapLinearFilter;
    texture.magFilter = LinearFilter;
    texture.generateMipmaps = false;

    texture.needsUpdate = true;
    // TODO: Detect NoColorSpace vs. LinearSRGBColorSpace based on primaries.
    texture.colorSpace = dfdTransferFn === KHR_DF_TRANSFER_SRGB ? SRGBColorSpace : NoColorSpace;
    texture.premultiplyAlpha = !!(dfdFlags & KHR_DF_FLAG_ALPHA_PREMULTIPLIED);

    return texture;
  }

  /**
   * @param {ArrayBuffer} buffer
   * @param {object?} config
   * @return {Promise<CompressedTexture|CompressedArrayTexture|DataTexture|Data3DTexture>}
   */
  async _createTexture(buffer, config = {}) {
    const container = read(new Uint8Array(buffer));

    if (container.vkFormat !== VK_FORMAT_UNDEFINED) {
      return createDataTexture(container);
    }

    //
    const taskConfig = config;
    const texturePending = this.init()
      .then(() => {
        return this.workerPool.postMessage({ type: "transcode", buffer, taskConfig: taskConfig }, [buffer]);
      })
      .then((e) => this._createTextureFrom(e.data, container));

    // Cache the task result.
    _taskCache.set(buffer, { promise: texturePending });

    return texturePending;
  }

  dispose() {
    this.workerPool.dispose();
    if (this.workerSourceURL) URL.revokeObjectURL(this.workerSourceURL);

    _activeLoaders--;

    return this;
  }
}

/* CONSTANTS */

KTX2Loader.BasisFormat = {
  ETC1S: 0,
  UASTC_4x4: 1,
};

KTX2Loader.TranscoderFormat = {
  ETC1: 0,
  ETC2: 1,
  BC1: 2,
  BC3: 3,
  BC4: 4,
  BC5: 5,
  BC7_M6_OPAQUE_ONLY: 6,
  BC7_M5: 7,
  PVRTC1_4_RGB: 8,
  PVRTC1_4_RGBA: 9,
  ASTC_4x4: 10,
  ATC_RGB: 11,
  ATC_RGBA_INTERPOLATED_ALPHA: 12,
  RGBA32: 13,
  RGB565: 14,
  BGR565: 15,
  RGBA4444: 16,
};

KTX2Loader.EngineFormat = {
  RGBAFormat: RGBAFormat,
  RGBA_ASTC_4x4_Format: RGBA_ASTC_4x4_Format,
  RGBA_BPTC_Format: RGBA_BPTC_Format,
  RGBA_ETC2_EAC_Format: RGBA_ETC2_EAC_Format,
  RGBA_PVRTC_4BPPV1_Format: RGBA_PVRTC_4BPPV1_Format,
  RGBA_S3TC_DXT5_Format: RGBA_S3TC_DXT5_Format,
  RGB_ETC1_Format: RGB_ETC1_Format,
  RGB_ETC2_Format: RGB_ETC2_Format,
  RGB_PVRTC_4BPPV1_Format: RGB_PVRTC_4BPPV1_Format,
  RGB_S3TC_DXT1_Format: RGB_S3TC_DXT1_Format,
};

//
// DataTexture and Data3DTexture parsing.

const FORMAT_MAP = {
  [VK_FORMAT_R32G32B32A32_SFLOAT]: RGBAFormat,
  [VK_FORMAT_R16G16B16A16_SFLOAT]: RGBAFormat,
  [VK_FORMAT_R8G8B8A8_UNORM]: RGBAFormat,
  [VK_FORMAT_R8G8B8A8_SRGB]: RGBAFormat,

  [VK_FORMAT_R32G32_SFLOAT]: RGFormat,
  [VK_FORMAT_R16G16_SFLOAT]: RGFormat,
  [VK_FORMAT_R8G8_UNORM]: RGFormat,
  [VK_FORMAT_R8G8_SRGB]: RGFormat,

  [VK_FORMAT_R32_SFLOAT]: RedFormat,
  [VK_FORMAT_R16_SFLOAT]: RedFormat,
  [VK_FORMAT_R8_SRGB]: RedFormat,
  [VK_FORMAT_R8_UNORM]: RedFormat,
};

const TYPE_MAP = {
  [VK_FORMAT_R32G32B32A32_SFLOAT]: FloatType,
  [VK_FORMAT_R16G16B16A16_SFLOAT]: HalfFloatType,
  [VK_FORMAT_R8G8B8A8_UNORM]: UnsignedByteType,
  [VK_FORMAT_R8G8B8A8_SRGB]: UnsignedByteType,

  [VK_FORMAT_R32G32_SFLOAT]: FloatType,
  [VK_FORMAT_R16G16_SFLOAT]: HalfFloatType,
  [VK_FORMAT_R8G8_UNORM]: UnsignedByteType,
  [VK_FORMAT_R8G8_SRGB]: UnsignedByteType,

  [VK_FORMAT_R32_SFLOAT]: FloatType,
  [VK_FORMAT_R16_SFLOAT]: HalfFloatType,
  [VK_FORMAT_R8_SRGB]: UnsignedByteType,
  [VK_FORMAT_R8_UNORM]: UnsignedByteType,
};

const COLOR_SPACE_MAP = {
  [VK_FORMAT_R8G8B8A8_SRGB]: SRGBColorSpace,
  [VK_FORMAT_R8G8_SRGB]: SRGBColorSpace,
  [VK_FORMAT_R8_SRGB]: SRGBColorSpace,
};

async function createDataTexture(container) {
  const { vkFormat, pixelWidth, pixelHeight, pixelDepth } = container;

  if (FORMAT_MAP[vkFormat] === undefined) {
    throw new Error("THREE.KTX2Loader: Unsupported vkFormat.");
  }

  const level = container.levels[0];

  let levelData;
  let view;

  if (container.supercompressionScheme === KHR_SUPERCOMPRESSION_NONE) {
    levelData = level.levelData;
  } else if (container.supercompressionScheme === KHR_SUPERCOMPRESSION_ZSTD) {
    if (!_zstd) {
      _zstd = new Promise(async (resolve) => {
        const zstd = new ZSTDDecoder();
        await zstd.init();
        resolve(zstd);
      });
    }

    levelData = (await _zstd).decode(level.levelData, level.uncompressedByteLength);
  } else {
    throw new Error("THREE.KTX2Loader: Unsupported supercompressionScheme.");
  }

  if (TYPE_MAP[vkFormat] === FloatType) {
    view = new Float32Array(levelData.buffer, levelData.byteOffset, levelData.byteLength / Float32Array.BYTES_PER_ELEMENT);
  } else if (TYPE_MAP[vkFormat] === HalfFloatType) {
    view = new Uint16Array(levelData.buffer, levelData.byteOffset, levelData.byteLength / Uint16Array.BYTES_PER_ELEMENT);
  } else {
    view = levelData;
  }
  //

  const texture =
    pixelDepth === 0
      ? new DataTexture(view, pixelWidth, pixelHeight)
      : new Data3DTexture(view, pixelWidth, pixelHeight, pixelDepth);

  texture.type = TYPE_MAP[vkFormat];
  texture.format = FORMAT_MAP[vkFormat];
  texture.colorSpace = COLOR_SPACE_MAP[vkFormat] || NoColorSpace;

  texture.needsUpdate = true;

  //

  return Promise.resolve(texture);
}

export { KTX2Loader };
